from pyspark.ml.feature import Word2Vec
from pyspark.ml.feature import Word2VecModel
from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark import RDD
import os
from pymongo import MongoClient
import time
import Config as conf

class PipeUtil:

    def __init__(self, rootPath):
        self.rootPath = rootPath
        self.spark = SparkSession.builder.appName('appName').master('local').config('spark.driver.cores', conf.SPARK_CORE).config("spark.driver.memory", "4g").config("spark.executor.memory", "4g").getOrCreate()
        self.sc = self.spark.sparkContext
        # self.dbClient = MongoClient()
        # self.db = self.dbClient['dbPipe']
        # self.collection = self.db[collectionName]
        self.model = self.__model__(self.rootPath)
        self.dataDict = {}

    def __model__(self, rootPath):
        path = rootPath + 'model'

        if os.path.exists(path):
            model = Word2VecModel.load(path)
        else:
            # 默认使用第一天的数据作为样本空间来简化运算
            rdd = self.sc.textFile(self.rootPath + 'date=01/part-*')
            parts = rdd.map(lambda l: l.split())
            lines = parts.map(lambda p: Row(content=p[1].split(',')))

            documentDF = self.spark.createDataFrame(lines)

            # Learn a mapping from words to Vectors.
            word2Vec = Word2Vec(vectorSize=3, minCount=0, inputCol="content", outputCol="result")
            model = word2Vec.fit(documentDF)

            # save model
            Word2VecModel.save(model, path=path)
        return model

    # def __del__(self):
    #     self.dbClient.close()
    #
    # def __mongoSave__(self, record):
    #     self.collection.insert(record)
    #
    # def __mongoLoad__(self, userId, clickTime):
    #     newsTotal = ''
    #
    #     for record in self.collection.find({'user_id': userId, 'time': {'$lte': int(clickTime)}}):
    #         newsTotal += record['content'] + ','
    #
    #     newsTotal = newsTotal[:-1]
    #
    #     return newsTotal

    # use memory
    def __newsByMemory__(self, rootPath, userId, clickTime):
        if not os.path.exists(rootPath):
            return ''

        dataDict = {}
        # serializationPath = rootPath + 'pythonData'

        if rootPath not in self.dataDict:
            self.dataDict.clear()
            i = 0
            while True:
                if i < 10:
                    subPath = 'part-0000' + str(i)
                else:
                    subPath = 'part-000' + str(i)

                file = rootPath + subPath
                if os.path.exists(file):
                    print(file)
                    with open(file, encoding='utf-8') as f:
                        for line in f:
                            tmp = line.split()
                            if len(tmp) == 0:
                                continue
                            else:
                                try:
                                    user_id = tmp[0]
                                    content = tmp[1]
                                    timeInt = int(tmp[2])

                                    data = {'user_id': tmp[0], 'content': tmp[1], 'time': int(tmp[2])}

                                    if user_id in dataDict:
                                        if timeInt <= int(clickTime):
                                            dataDict[user_id]['content'] += ',' + content
                                    else:
                                        dataDict[user_id] = data
                                except:
                                    print('bad data:{}'.format(tmp))
                    i += 1
                else:
                    break

            self.dataDict[rootPath] = dataDict

            # with open(serializationPath, 'wb') as file:
            #     print('start dump file')
            #     dumpStart = time.time()
            #     pickle.dump(dataList, file, True)
            #     print('dump file time:{}'.format(time.time() - dumpStart))
        else:
            # with open(serializationPath, 'rb') as file:
            #     print('start load file')
            #     loadStart = time.time()
            #     dataList = pickle.load(file)
            #     print('load file time:{}'.format(time.time() - loadStart))

            dataDict = self.dataDict[rootPath]

        # # filter func
        # def filterFunc(data):
        #     return data['user_id'] == userId and data['time'] <= int(clickTime)
        # print('filter data')
        # datas = filter(filterFunc, dataList)
        # newsTotal = ''
        # for record in datas:
        #     newsTotal += record['content'] + ','
        #
        # newsTotal = newsTotal[:-1]
        if userId in dataDict:
            return dataDict[userId]['content']
        else:
            return ''

    # use spark
    def __newsBySpark__(self, rootPath, userId, clickTime):
        objectPath = rootPath + '/object'
        if os.path.exists(objectPath):
            objectStart = time.time()
            print('loading rdd object')
            rddUser = self.sc.pickleFile(objectPath)
            print('load rdd object time:{}'.format(time.time() - objectStart))
        else:
            fileStart = time.time()
            print('reading rdd object')
            rddUser = self.sc.textFile(rootPath + '/part-*')
            print('read rdd object time:{}'.format(time.time() - fileStart))

            saveStart = time.time()
            print('saving rdd file')
            RDD.saveAsPickleFile(rddUser, objectPath)
            print('save rdd file time:{}'.format(time.time() - saveStart))

        print('start parts')
        parts = rddUser.map(lambda l: l.split())
        print('start lines')
        lines = parts.map(lambda p: Row(user_id=p[0], content=p[1], time=int(p[2])))
        print('start createDataFrame')
        df = self.spark.createDataFrame(lines)
        print('start createOrReplaceTempView')
        df.createOrReplaceTempView("table_news")

        print('start sql 0')
        users = self.spark.sql(
            "SELECT content FROM table_news WHERE user_id == " + userId + " AND time <= " + str(clickTime))
        newsTotal = ''

        print('start collect')
        collectStart = time.time()
        self.contentRdd = users.rdd
        self.contentRdd.persist()
        newsCollect = self.contentRdd.collect()
        print('collect time:{}'.format(time.time() - collectStart))
        for content in newsCollect:
            newsTotal += content.content + ','
        # remove last ','
        newsTotal = newsTotal[:-1]
        print(newsTotal)
        return newsTotal

    # use mongodb
    # def __newsByDB__(self, rootPath, userId, clickTime):
    #     dbPath = rootPath + 'db'
    #     if not os.path.exists(dbPath):
    #         for parent, dirnames, filenames in os.walk(rootPath):
    #             for filename in filenames:
    #                 if not filename.startswith('part-'):
    #                     continue
    #                 else:
    #                     file = os.path.join(parent, filename)
    #                     print(file)
    #                     with open(file, encoding='utf-8') as f:
    #                         for line in f:
    #                             tmp = line.split()
    #                             if len(tmp) == 0:
    #                                 continue
    #                             else:
    #                                 data = {'user_id': tmp[0], 'content': tmp[1], 'time': int(tmp[2])}
    #                                 self.__mongoSave__(data)
    #         # index
    #         self.collection.create_index([('user_id', 1), ('time', 1)])
    #
    #         # 占位文件
    #         with open(dbPath, 'w') as f:
    #             f.write('already saved in mongodb')
    #
    #     return self.__mongoLoad__(userId, clickTime)

    def vector(self, userId, clickTime, isPredict = False):

        # concatenate
        day = clickTime[0:2]
        dayInt = int(day)
        count = 1 # 默认1天
        if isPredict:
            count = 3 # 考虑3天的数据,predict

        def dayName(d):
            if d < 0:
                return ''
            elif 0 <= d < 10:
                return '0' + str(d)
            else:
                return str(d)

        newsTotal = ''
        for i in range(count):
            dayIndex = dayName(dayInt - i)
            newsStr = self.__newsByMemory__(self.rootPath + 'date=' + dayIndex + '/', userId, clickTime)
            if newsStr:

                if newsTotal:
                    newsTotal += ','

                newsTotal += newsStr

        return newsTotal

        # # print('newsStr:{}'.format(newsTotal))
        # # 获取对应vector
        # print('create data frame')
        # testDF = self.spark.createDataFrame([(newsTotal.split(','),)], ['content'])
        # print('transform data')
        # result = self.model.transform(testDF).head().result
        # return result
